#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
企业舆情分析系统
整合多AI对话和问题生成功能，提供完整的企业分析工作流程
"""

import asyncio
import aiohttp
import sqlite3
import json
import time
import uuid
from datetime import datetime
from typing import Dict, List, Optional, Any, Tuple
from dataclasses import dataclass, field

# 导入我们的模块
from multi_ai_chat import (
    create_multi_ai_chat, 
    setup_example_providers, 
    AIResponse,
    DeepSeekProvider as MultiAIDeepSeekProvider
)
from question_generator import (
    create_question_generator,
    Company,
    QuestionType,
    GeneratedQuestion
)


@dataclass
class EnterpriseInfo:
    """企业信息类"""
    name: str
    industry: str = ""
    description: str = ""
    competitors: List[str] = field(default_factory=list)
    keywords: List[str] = field(default_factory=list)
    founded_year: Optional[int] = None
    headquarters: str = ""
    market_cap: str = ""
    business_scope: str = ""
    risk_level: str = "low"
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class AnalysisResult:
    """分析结果类"""
    session_id: str
    enterprise_name: str
    question: str
    question_type: str
    ai_responses: List[AIResponse]
    timestamp: datetime
    summary: Dict[str, Any]
    metadata: Dict[str, Any] = field(default_factory=dict)


class EnterpriseInfoExtractor:
    """企业信息提取器（使用DeepSeek）"""
    
    def __init__(self, api_key: str):
        self.api_key = api_key
        self.base_url = "https://api.deepseek.com/v1"
    
    async def extract_enterprise_info(self, enterprise_name: str) -> EnterpriseInfo:
        """使用DeepSeek提取企业信息"""
        timeout = aiohttp.ClientTimeout(total=30)
        async with aiohttp.ClientSession(timeout=timeout) as session:
            try:
                url = f"{self.base_url}/chat/completions"
                headers = {
                    'Authorization': f'Bearer {self.api_key}',
                    'Content-Type': 'application/json'
                }
                
                system_prompt = """你是一个专业的企业信息分析专家。请根据给定的企业名称，提供详细的企业信息。

请按照以下JSON格式返回信息（只返回JSON，不要其他内容）：
{
    "name": "企业名称",
    "industry": "所属行业",
    "description": "企业简介（100字以内）",
    "competitors": ["竞争对手1", "竞争对手2", "竞争对手3"],
    "keywords": ["关键词1", "关键词2", "关键词3", "关键词4"],
    "founded_year": 创立年份(数字),
    "headquarters": "总部位置",
    "market_cap": "市值规模（如：千亿级、百亿级等）",
    "business_scope": "主要业务范围",
    "risk_level": "风险等级（低/中/高）"
}

如果企业信息不确定，请在相应字段填入"未知"或空数组。"""

                payload = {
                    "model": "deepseek-chat",
                    "messages": [
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": f"请分析企业：{enterprise_name}"}
                    ],
                    "temperature": 0.3,
                    "max_tokens": 1000
                }
                
                async with session.post(url, json=payload, headers=headers) as response:
                    if response.status == 200:
                        data = await response.json()
                        if 'choices' in data and len(data['choices']) > 0:
                            content = data['choices'][0]['message']['content'].strip()
                            
                            # 尝试解析JSON
                            try:
                                # 提取JSON部分
                                start_idx = content.find('{')
                                end_idx = content.rfind('}') + 1
                                if start_idx != -1 and end_idx != -1:
                                    json_str = content[start_idx:end_idx]
                                    info_dict = json.loads(json_str)
                                    
                                    return EnterpriseInfo(
                                        name=info_dict.get('name', enterprise_name),
                                        industry=info_dict.get('industry', ''),
                                        description=info_dict.get('description', ''),
                                        competitors=info_dict.get('competitors', []),
                                        keywords=info_dict.get('keywords', [enterprise_name]),
                                        founded_year=info_dict.get('founded_year'),
                                        headquarters=info_dict.get('headquarters', ''),
                                        market_cap=info_dict.get('market_cap', ''),
                                        business_scope=info_dict.get('business_scope', ''),
                                        risk_level=info_dict.get('risk_level', ''),
                                        metadata={"source": "deepseek", "raw_response": content}
                                    )
                            except json.JSONDecodeError:
                                print(f"⚠️ JSON解析失败，使用基本信息")
                                
                    else:
                        error_text = await response.text()
                        print(f"DeepSeek API错误: {response.status} - {error_text}")
                        
            except Exception as e:
                print(f"企业信息提取异常: {e}")
        
        # 如果提取失败，返回基本信息
        return EnterpriseInfo(
            name=enterprise_name,
            keywords=[enterprise_name],
            metadata={"source": "fallback"}
        )


class DatabaseManager:
    """数据库管理器"""
    
    def __init__(self, db_path: str = "enterprise_analysis.db"):
        self.db_path = db_path
        self.init_database()
    
    def init_database(self):
        """初始化数据库表"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 企业信息表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS enterprises (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT UNIQUE NOT NULL,
                industry TEXT,
                description TEXT,
                competitors TEXT,
                keywords TEXT,
                founded_year INTEGER,
                headquarters TEXT,
                market_cap TEXT,
                business_scope TEXT,
                risk_level TEXT,
                metadata TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # 分析会话表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS analysis_sessions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                session_id TEXT UNIQUE NOT NULL,
                enterprise_name TEXT NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                completed_at TIMESTAMP,
                status TEXT DEFAULT 'running',
                total_questions INTEGER DEFAULT 0,
                successful_responses INTEGER DEFAULT 0,
                metadata TEXT
            )
        ''')
        
        # 问题记录表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS questions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                session_id TEXT NOT NULL,
                question_text TEXT NOT NULL,
                question_type TEXT NOT NULL,
                is_inspirational BOOLEAN DEFAULT FALSE,
                generated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                metadata TEXT
            )
        ''')
        
        # AI响应表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS ai_responses (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                session_id TEXT NOT NULL,
                question_id INTEGER NOT NULL,
                ai_provider TEXT NOT NULL,
                response_text TEXT,
                success BOOLEAN DEFAULT TRUE,
                response_time REAL,
                error_message TEXT,
                usage_info TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (question_id) REFERENCES questions (id)
            )
        ''')
        
        conn.commit()
        conn.close()
    
    def save_enterprise_info(self, enterprise_info: EnterpriseInfo) -> int:
        """保存企业信息"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        try:
            cursor.execute('''
                INSERT OR REPLACE INTO enterprises 
                (name, industry, description, competitors, keywords, founded_year, 
                 headquarters, market_cap, business_scope, risk_level, metadata, updated_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
            ''', (
                enterprise_info.name,
                enterprise_info.industry,
                enterprise_info.description,
                json.dumps(enterprise_info.competitors, ensure_ascii=False),
                json.dumps(enterprise_info.keywords, ensure_ascii=False),
                enterprise_info.founded_year,
                enterprise_info.headquarters,
                enterprise_info.market_cap,
                enterprise_info.business_scope,
                enterprise_info.risk_level,
                json.dumps(enterprise_info.metadata, ensure_ascii=False)
            ))
            
            enterprise_id = cursor.lastrowid
            conn.commit()
            return enterprise_id
            
        finally:
            conn.close()
    
    def create_analysis_session(self, enterprise_name: str, metadata: Dict = None) -> str:
        """创建分析会话"""
        session_id = str(uuid.uuid4())
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        try:
            cursor.execute('''
                INSERT INTO analysis_sessions (session_id, enterprise_name, metadata)
                VALUES (?, ?, ?)
            ''', (
                session_id,
                enterprise_name,
                json.dumps(metadata or {}, ensure_ascii=False)
            ))
            
            conn.commit()
            return session_id
            
        finally:
            conn.close()
    
    def save_question(self, session_id: str, question: GeneratedQuestion) -> int:
        """保存问题"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        try:
            cursor.execute('''
                INSERT INTO questions (session_id, question_text, question_type, is_inspirational, metadata)
                VALUES (?, ?, ?, ?, ?)
            ''', (
                session_id,
                question.question,
                question.question_type.value,
                question.is_inspirational,
                json.dumps(question.metadata, ensure_ascii=False)
            ))
            
            question_id = cursor.lastrowid
            conn.commit()
            return question_id
            
        finally:
            conn.close()
    
    def save_ai_responses(self, session_id: str, question_id: int, responses: List[AIResponse]):
        """保存AI响应"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        try:
            for response in responses:
                cursor.execute('''
                    INSERT INTO ai_responses 
                    (session_id, question_id, ai_provider, response_text, success, 
                     response_time, error_message, usage_info)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                ''', (
                    session_id,
                    question_id,
                    response.provider,
                    response.response if response.success else None,
                    response.success,
                    response.response_time,
                    response.error if not response.success else None,
                    json.dumps(response.metadata, ensure_ascii=False)
                ))
            
            conn.commit()
            
        finally:
            conn.close()
    
    def update_session_status(self, session_id: str, status: str, total_questions: int = None, 
                            successful_responses: int = None):
        """更新会话状态"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        try:
            updates = ["status = ?"]
            params = [status]
            
            if total_questions is not None:
                updates.append("total_questions = ?")
                params.append(total_questions)
            
            if successful_responses is not None:
                updates.append("successful_responses = ?")
                params.append(successful_responses)
            
            if status == 'completed':
                updates.append("completed_at = CURRENT_TIMESTAMP")
            
            params.append(session_id)
            
            cursor.execute(f'''
                UPDATE analysis_sessions 
                SET {", ".join(updates)}
                WHERE session_id = ?
            ''', params)
            
            conn.commit()
            
        finally:
            conn.close()
    
    def get_enterprise_analysis_history(self, enterprise_name: str, limit: int = 10) -> List[Dict]:
        """获取企业分析历史"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        try:
            cursor.execute('''
                SELECT session_id, created_at, completed_at, status, total_questions, successful_responses
                FROM analysis_sessions 
                WHERE enterprise_name = ?
                ORDER BY created_at DESC
                LIMIT ?
            ''', (enterprise_name, limit))
            
            columns = [description[0] for description in cursor.description]
            results = []
            for row in cursor.fetchall():
                results.append(dict(zip(columns, row)))
            
            return results
            
        finally:
            conn.close()
    
    def get_session_details(self, session_id: str) -> Dict[str, Any]:
        """获取会话详细信息"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        try:
            # 获取会话信息
            cursor.execute('''
                SELECT * FROM analysis_sessions WHERE session_id = ?
            ''', (session_id,))
            session_data = cursor.fetchone()
            
            if not session_data:
                return {}
            
            # 获取问题和响应
            cursor.execute('''
                SELECT q.id, q.question_text, q.question_type, q.is_inspirational,
                       r.ai_provider, r.response_text, r.success, r.response_time, r.error_message
                FROM questions q
                LEFT JOIN ai_responses r ON q.id = r.question_id
                WHERE q.session_id = ?
                ORDER BY q.id, r.ai_provider
            ''', (session_id,))
            
            questions_data = cursor.fetchall()
            
            return {
                "session": session_data,
                "questions_and_responses": questions_data
            }
            
        finally:
            conn.close()


class EnterpriseAnalysisSystem:
    """企业舆情分析系统主类"""
    
    def __init__(self, config: Dict[str, Dict[str, str]], deepseek_api_key: str = None):
        """
        初始化系统
        
        Args:
            config: AI服务配置字典
            deepseek_api_key: DeepSeek API密钥（用于企业信息提取）
        """
        self.config = config
        self.deepseek_api_key = deepseek_api_key
        
        # 初始化组件
        self.multi_ai_chat = create_multi_ai_chat()
        setup_example_providers(self.multi_ai_chat, config)
        
        self.question_generator = create_question_generator(deepseek_api_key)
        
        if deepseek_api_key:
            self.info_extractor = EnterpriseInfoExtractor(deepseek_api_key)
        else:
            self.info_extractor = None
        
        self.db_manager = DatabaseManager()
    
    async def analyze_enterprise(self, enterprise_name: str, 
                               question_types: List[QuestionType] = None,
                               use_inspiration: bool = True,
                               max_questions_per_type: int = 3,
                               min_total_questions: int = 30) -> str:
        """
        完整的企业分析流程
        
        Args:
            enterprise_name: 企业名称
            question_types: 要分析的问题类型列表，None表示使用默认类型
            use_inspiration: 是否使用启发式问题生成
            max_questions_per_type: 每种类型最多生成几个问题
            
        Returns:
            session_id: 分析会话ID
        """
        print(f"🚀 开始分析企业: {enterprise_name}")
        
        # 1. 创建分析会话
        session_id = self.db_manager.create_analysis_session(
            enterprise_name, 
            {
                "use_inspiration": use_inspiration, 
                "max_questions_per_type": max_questions_per_type,
                "min_total_questions": min_total_questions
            }
        )
        
        try:
            # 2. 提取企业信息
            print("📊 正在提取企业信息...")
            if self.info_extractor:
                enterprise_info = await self.info_extractor.extract_enterprise_info(enterprise_name)
                print(f"✅ 企业信息提取完成: {enterprise_info.industry}")
            else:
                enterprise_info = EnterpriseInfo(name=enterprise_name)
                print("⚠️ 未配置DeepSeek API，使用基本企业信息")
            
            # 保存企业信息
            self.db_manager.save_enterprise_info(enterprise_info)
            
            # 3. 将企业信息添加到问题生成器
            company = Company(
                name=enterprise_info.name,
                industry=enterprise_info.industry,
                description=enterprise_info.description,
                competitors=enterprise_info.competitors,
                keywords=enterprise_info.keywords,
                founded_year=enterprise_info.founded_year,
                headquarters=enterprise_info.headquarters
            )
            self.question_generator.add_company(company)
            
            # 4. 生成问题 - 扩展问题范围以进行更全面的极端情况检测
            print("❓ 正在生成分析问题...")
            if question_types is None:
                # 使用更全面的问题类型组合进行极端情况检测
                if enterprise_info.risk_level == "高":
                    question_types = [
                        QuestionType.NEGATIVE_NEWS, QuestionType.REGULATION, QuestionType.FINANCIAL_INFO,
                        QuestionType.REPUTATION, QuestionType.LEADERSHIP, QuestionType.ESG,
                        QuestionType.COMPETITIVE_ANALYSIS, QuestionType.INVESTMENT
                    ]
                else:
                    # 即使是低风险企业，也要进行全面检测
                    question_types = [
                        QuestionType.NEGATIVE_NEWS, QuestionType.FINANCIAL_INFO, QuestionType.MARKET_ANALYSIS,
                        QuestionType.REPUTATION, QuestionType.LEADERSHIP, QuestionType.TECHNOLOGY,
                        QuestionType.COMPETITIVE_ANALYSIS, QuestionType.REGULATION, QuestionType.ESG
                    ]
            
            print(f"🎯 极端情况检测策略：覆盖 {len(question_types)} 种问题类型，目标生成至少 {min_total_questions} 个问题")
            print(f"📊 问题类型: {[qt.value for qt in question_types]}")
            
            questions = await self.question_generator.generate_batch_questions(
                [enterprise_name],
                question_types,
                use_inspiration=use_inspiration,
                max_per_company=max_questions_per_type,
                min_total_questions=min_total_questions
            )
            
            print(f"✅ 生成了 {len(questions)} 个问题")
            
            # 5. 向多个AI发送问题并保存结果
            total_successful = 0
            for i, question in enumerate(questions, 1):
                print(f"🤖 正在查询问题 {i}/{len(questions)}: {question.question[:50]}...")
                
                # 保存问题
                question_id = self.db_manager.save_question(session_id, question)
                
                # 向多个AI发送问题
                responses = await self.multi_ai_chat.send_to_all(question.question)
                
                # 保存响应
                self.db_manager.save_ai_responses(session_id, question_id, responses)
                
                # 统计成功响应
                successful_count = sum(1 for r in responses if r.success)
                total_successful += successful_count
                
                print(f"   📝 收到 {successful_count}/{len(responses)} 个成功响应")
            
            # 6. 更新会话状态
            self.db_manager.update_session_status(
                session_id, 
                'completed',
                len(questions),
                total_successful
            )
            
            print(f"🎉 企业分析完成！会话ID: {session_id}")
            print(f"📈 总计: {len(questions)} 个问题, {total_successful} 个成功响应")
            
            return session_id
            
        except Exception as e:
            print(f"❌ 分析过程中发生错误: {e}")
            self.db_manager.update_session_status(session_id, 'error')
            raise
    
    def get_analysis_summary(self, session_id: str) -> Dict[str, Any]:
        """获取分析结果摘要"""
        details = self.db_manager.get_session_details(session_id)
        
        if not details:
            return {"error": "会话不存在"}
        
        session_info = details["session"]
        questions_data = details["questions_and_responses"]
        
        # 整理数据
        questions_summary = {}
        ai_performance = {}
        
        for row in questions_data:
            q_id, q_text, q_type, is_inspirational, ai_provider, response_text, success, response_time, error = row
            
            # 问题统计
            if q_id not in questions_summary:
                questions_summary[q_id] = {
                    "text": q_text,
                    "type": q_type,
                    "inspirational": bool(is_inspirational),
                    "responses": []
                }
            
            if ai_provider:
                questions_summary[q_id]["responses"].append({
                    "provider": ai_provider,
                    "success": bool(success),
                    "response": response_text,
                    "response_time": response_time,
                    "error": error
                })
                
                # AI性能统计
                if ai_provider not in ai_performance:
                    ai_performance[ai_provider] = {"success": 0, "total": 0, "avg_time": 0}
                
                ai_performance[ai_provider]["total"] += 1
                if success:
                    ai_performance[ai_provider]["success"] += 1
                if response_time:
                    ai_performance[ai_provider]["avg_time"] += response_time
        
        # 计算平均响应时间
        for provider in ai_performance:
            total = ai_performance[provider]["total"]
            if total > 0:
                ai_performance[provider]["avg_time"] /= total
                ai_performance[provider]["success_rate"] = ai_performance[provider]["success"] / total
        
        return {
            "session_id": session_id,
            "enterprise_name": session_info[2],
            "status": session_info[5],
            "created_at": session_info[3],
            "completed_at": session_info[4],
            "total_questions": len(questions_summary),
            "ai_performance": ai_performance,
            "questions": list(questions_summary.values())
        }
    
    def print_analysis_result(self, session_id: str):
        """打印分析结果"""
        summary = self.get_analysis_summary(session_id)
        
        if "error" in summary:
            print(f"❌ {summary['error']}")
            return
        
        print(f"\n{'='*80}")
        print(f"📊 企业舆情分析报告")
        print(f"{'='*80}")
        print(f"🏢 企业名称: {summary['enterprise_name']}")
        print(f"🆔 会话ID: {summary['session_id']}")
        print(f"📅 分析时间: {summary['created_at']}")
        print(f"📈 问题总数: {summary['total_questions']}")
        print(f"🎯 分析状态: {summary['status']}")
        
        # AI性能统计
        print(f"\n🤖 AI服务性能统计:")
        print("-" * 60)
        for provider, perf in summary['ai_performance'].items():
            print(f"  {provider:15} | 成功率: {perf['success_rate']:.1%} | "
                  f"响应时间: {perf['avg_time']:.2f}s | "
                  f"总请求: {perf['total']}")
        
        # 问题和回答概览
        print(f"\n📋 问题和回答概览:")
        print("-" * 60)
        for i, question in enumerate(summary['questions'][:5], 1):  # 只显示前5个
            inspiration_mark = "✨" if question['inspirational'] else "📝"
            print(f"\n{i}. {inspiration_mark} [{question['type']}]")
            print(f"   问题: {question['text']}")
            
            successful_responses = [r for r in question['responses'] if r['success']]
            print(f"   成功响应: {len(successful_responses)}/{len(question['responses'])}")
            
            if successful_responses:
                # 显示最短的一个回答作为示例
                shortest_response = min(successful_responses, key=lambda x: len(x['response'] or ''))
                response_preview = (shortest_response['response'] or '')[:100]
                print(f"   示例回答({shortest_response['provider']}): {response_preview}...")
        
        if len(summary['questions']) > 5:
            print(f"\n   ... 还有 {len(summary['questions']) - 5} 个问题（完整结果请查看数据库）")
    
    def export_analysis_report(self, session_id: str, format: str = "json") -> str:
        """导出分析报告"""
        summary = self.get_analysis_summary(session_id)
        
        if format.lower() == "json":
            return json.dumps(summary, ensure_ascii=False, indent=2)
        
        elif format.lower() == "txt":
            lines = [
                f"企业舆情分析报告",
                f"=" * 50,
                f"企业名称: {summary['enterprise_name']}",
                f"分析时间: {summary['created_at']}",
                f"问题总数: {summary['total_questions']}",
                f"状态: {summary['status']}",
                "",
                "AI服务性能:",
                "-" * 30
            ]
            
            for provider, perf in summary['ai_performance'].items():
                lines.append(f"{provider}: 成功率{perf['success_rate']:.1%}, "
                           f"平均响应时间{perf['avg_time']:.2f}s")
            
            lines.extend(["", "问题详情:", "-" * 30])
            
            for i, question in enumerate(summary['questions'], 1):
                inspiration_mark = "✨" if question['inspirational'] else "📝"
                lines.append(f"{i}. {inspiration_mark} [{question['type']}]")
                lines.append(f"   {question['text']}")
                
                for response in question['responses']:
                    if response['success']:
                        lines.append(f"   {response['provider']}: {response['response'][:200]}...")
                    else:
                        lines.append(f"   {response['provider']}: 错误 - {response['error']}")
                lines.append("")
            
            return "\n".join(lines)
        
        else:
            return str(summary)


# 便捷函数
def create_analysis_system(config_file: str = "ai_config.json", 
                          deepseek_api_key: str = None) -> EnterpriseAnalysisSystem:
    """创建企业分析系统实例"""
    
    # 加载配置
    try:
        with open(config_file, 'r', encoding='utf-8') as f:
            config = json.load(f)
    except FileNotFoundError:
        print(f"⚠️ 配置文件 {config_file} 不存在")
        config = {}
    except json.JSONDecodeError:
        print(f"⚠️ 配置文件 {config_file} 格式错误")
        config = {}
    
    # 如果没有单独提供DeepSeek密钥，尝试从配置文件获取
    if not deepseek_api_key and "deepseek" in config:
        deepseek_api_key = config["deepseek"]["api_key"]
    
    return EnterpriseAnalysisSystem(config, deepseek_api_key)


if __name__ == "__main__":
    print("🎯 企业舆情分析系统已加载完成！")
    print("💡 使用说明:")
    print("   1. 创建系统: system = create_analysis_system()")
    print("   2. 分析企业: await system.analyze_enterprise('企业名称')")
    print("   3. 数据存储在enterprise_analysis.db中") 